from llama_index.core.agent import ReActAgent
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.llms.dashscope import DashScope
from llama_index.core.tools import FunctionTool

# 工具定义
def calculate(expression: str) -> str:
    """执行科学计算"""
    try:
        result = eval(expression)
        return str(result)
    except Exception as e:
        return f"计算错误: {e}"

def ask_fruit_unit_price(fruit: str) -> str:
    """查询水果价格"""
    fruit = fruit.lower()
    if '苹果' in fruit:
        return "10"
    elif '香蕉' in fruit:
        return "20"
    else:
        return "15"

# 创建工具
tools = [
    FunctionTool.from_defaults(
        fn=calculate,
        name="calculate",
        description="执行数学计算，输入数学表达式，返回计算结果"
    ),
    FunctionTool.from_defaults(
        fn=ask_fruit_unit_price,
        name="ask_fruit_unit_price",
        description="查询水果单价，输入水果名称，返回每千克价格"
    )
]

# 使用 Qwen 模型
llm = DashScope(model="qwen-plus", api_key="sk-965dc39b016c49ecbe29de180f4db2b6")
# 创建 memory
memory = ChatMemoryBuffer.from_defaults()
# 创建ReActAgent - 新版本直接传入参数
agent = ReActAgent(
    tools=tools,
    llm=llm,
    memory=memory,  # 添加 memory 参数
    verbose=True
)

# 执行查询
# response = agent.chat("请帮我计算3千克香蕉和2千克苹果的总价格。"
#     "请按照以下步骤执行："
#     "1. 首先使用 ask_fruit_unit_price 工具查询香蕉的单价"
#     "2. 然后使用 ask_fruit_unit_price 工具查询苹果的单价"
#     "3. 接着使用 calculate 工具计算香蕉的总价：香蕉单价 × 3"
#     "4. 再使用 calculate 工具计算苹果的总价：苹果单价 × 2"
#     "5. 最后使用 calculate 工具将两个总价相加"
#     "请确保调用相应的工具来完成计算。")
# print(response)


# 分步骤执行
print("步骤1: 查询香蕉价格")
price_banana_response = agent.chat("查询香蕉的单价")
print(price_banana_response)

print("\n步骤2: 查询苹果价格")
price_apple_response = agent.chat("查询苹果的单价")
print(price_apple_response)

print("\n步骤3: 计算总价")
final_response = agent.chat("基于前面的信息，计算3千克香蕉和2千克苹果的总价格")
print(final_response)